import os
import jittor
from jittor import transform
from ljp.dataset.cifar import CIFAR100
from ljp.metrics import accuracy
from model.resnet import resnet18
from ljp import Trainer
from model.resnet_sub import resnet18_sub, resnet50_sub

# os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 使用1号
if jittor.has_cuda:
    jittor.flags.use_cuda = 1

from jittor.dataset import MNIST


def get_train_transforms():
    return transform.Compose([
        transform.RandomCropAndResize((32, 32)),
        transform.RandomHorizontalFlip(),
        transform.ToTensor(),
        transform.ImageNormalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
    ])


def get_valid_transforms():
    return transform.Compose([
        transform.Resize(32),
        transform.CenterCrop(32),
        transform.ToTensor(),
        transform.ImageNormalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
    ])


batch_size = 128
train_loader = CIFAR100(train=True, batch_size=batch_size, shuffle=True, root=r'D:/data/cifar100/', download=False,
                        transform=get_train_transforms())
val_loader = CIFAR100(train=False, batch_size=batch_size, shuffle=True, root=r'D:/data/cifar100/', download=False,
                      transform=get_valid_transforms())

model = Trainer(resnet18(num_classes=100))
# model = Trainer(resnet50_sub(num_classes=100))

# model.compile(loss_func=jittor.nn.CrossEntropyLoss(), metrics_dict={"accuracy": accuracy})
loss_dict = {
    'loss_func': jittor.nn.CrossEntropyLoss,
}
optimizer_dict3 = {'optimizer': jittor.optim.SGD,
                   'lr': 1e-3,
                   }

lr_scheduler_dict = {
    'lr_scheduler': jittor.optim.LRScheduler,
    'step_size': 60,
    'gamma': 0.1,
}
model.compile(loss_dict=loss_dict,
              optimizer_dict=optimizer_dict3,
              # lr_scheduler_dict=lr_scheduler_dict,
              metrics_dict={"acc1": accuracy, },
              numpy_metric=False,
              monitor='val_acc1',
              monitor_mode='max')
model.fit(epochs=200, T_loader=train_loader, V_loader=val_loader)
